import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from target import DuffingOscillator
from flowDMD import FlowDMD, train_flowdmd_xy
from dictionary import DicRBF
from solver import KoopmanGeneralSolver
import argparse
import os
import torch


duffing_train = DuffingOscillator(n_init=1000, traj_len=10, seed=0)
data_x_train = duffing_train.generate_init_data()
data_y_train = duffing_train.generate_next_data(data_x_train)

duffing_valid = DuffingOscillator(n_init=300, traj_len=10, seed=0)
data_x_valid = duffing_valid.generate_init_data()
data_y_valid = duffing_valid.generate_next_data(data_x_valid)

data_x_train = data_x_train.float()
data_y_train = data_y_train.float()
data_train = [data_x_train, data_y_train]
data_valid = [data_x_valid, data_y_valid]


def trainmodel():
    parser = argparse.ArgumentParser(description='FlowDMD训练和预测')
    parser.add_argument('--train', default='True',
                        action='store_true', help='是否进行训练')
    parser.add_argument('--model_path', type=str,
                        default='model.pth', help='模型保存路径')
    parser.add_argument('--num_dictionaries', type=int,
                        default=22, help='字典函数数量')
    args = parser.parse_args()

    # 可视化数据
    fig, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(10, 5))
    ax1.plot(data_x_train[:, 0], data_x_train[:, 1], '.')

    ax1.set_xlabel(r'$x_{1}$')
    ax1.set_ylabel(r'$x_{2}$')
    ax2.plot(data_y_train[:, 0], data_y_train[:, 1], '.')
    ax2.set_xlabel(r'$x_{1}$')
    ax2.set_ylabel(r'$x_{2}$')

    # 创建模型
    model = FlowDMD(dim=2, num_layers=6, hidden_dim=64,
                    reg=1e-4, num_dictionaries=args.num_dictionaries)

    # 检查是否存在模型文件
    if os.path.exists(args.model_path) and not args.train:
        print(f"加载已有模型: {args.model_path}")
        model.load_state_dict(torch.load(args.model_path))
    else:
        print("开始训练模型...")
        k = train_flowdmd_xy(model, data_xy=data_train,
                             epochs=2000, lr=1e-3, alpha=1.0)
        # 保存模型
        print(k)
        torch.save(model.state_dict(), args.model_path)
        print(f"模型已保存到: {args.model_path}")

        # 打印每个字典的权重
        print("最终字典权重:", model.dictionary_weights.detach().numpy())
    return model, k


def RBFmodel():
    rbf_basis_func = DicRBF(rbf_number=100, regularizer=1e-1)
    rbf_basis_func.build(data_x_train)

    rbf_solver = KoopmanGeneralSolver(dic=rbf_basis_func,
                                      target_dim=2,
                                      reg=0.5)

    rbf_solver.build(data_train)
    fig, (ax1, ax2) = plt.subplots(
        1, 2, figsize=(10, 5), sharex=False, sharey=True)
    ax2.scatter(rbf_solver.eigenvalues.real, rbf_solver.eigenvalues.imag)
    ax2.set_xlabel(r'Re$(\mu)$')
    ax2.set_title('EDMD (RBF)')

    fig.tight_layout()
    return rbf_solver


def test():
    rbf_solver = RBFmodel()
    model = FlowDMD(dim=2, num_layers=6, hidden_dim=64,
                    reg=1e-4, num_dictionaries=22)
    model.load_state_dict(torch.load('model.pth'))
    traj_len = 50
    duffing_test = DuffingOscillator(n_init=1, traj_len=traj_len, seed=111)
    data_x_test = duffing_test.generate_init_data()
    data_y_test = duffing_test.generate_next_data(data_x_test)
    data_x_test = data_x_test.float()
    data_y_test = data_y_test.float()
    x_traj = data_x_test.numpy()

    x0_test = data_x_test[0]
    x0_test = x0_test.reshape(-1, x0_test.shape[-1])

    with torch.no_grad():
        x_est_traj_rbf = rbf_solver.predict(x0_test, traj_len).numpy()
        x_est_traj_DL = model.predict(x0_test, traj_len)[0].numpy()
    DL_error = np.sqrt(np.mean(np.square(x_est_traj_DL-x_traj)))
    rbf_error = np.sqrt(np.mean(np.square(x_est_traj_rbf-x_traj)))
    fig, (ax1, ax2) = plt.subplots(1, 2, sharey=False, figsize=(12, 5))

    t_grid = np.arange(0, traj_len)
    ax1.plot(t_grid, x_traj[:, 0], 'k-', label='Exact', markevery=2)
    ax1.plot(t_grid, x_est_traj_DL[:, 0], 'bo', label='EDMD-Flow', markevery=2)
    ax1.plot(t_grid, x_est_traj_rbf[:, 0], 'gs',
             label='EDMD (RBF)', markevery=2)

    ax2.plot(t_grid, x_traj[:, 1], 'k-', label='Exact', markevery=2)
    ax2.plot(t_grid, x_est_traj_DL[:, 1], 'bo', label='EDMD-Flow', markevery=2)
    ax2.plot(t_grid, x_est_traj_rbf[:, 1], 'gs',
             label='EDMD (RBF)', markevery=2)

    ax1.set_xlabel(r"$n$")
    ax1.set_ylabel(r"$x_{1}(n)$")
    ax1.legend(loc='best', prop={'size': 14})

    ax2.set_xlabel(r"$n$")
    ax2.set_ylabel(r"$x_{2}(n)$")
    ax2.legend(loc='best', prop={'size': 14})

    print("DL error: ", DL_error,
          "  RBF error: ", rbf_error)
    fig.tight_layout()
    fig.show()
    fig.savefig('test.png')


def test2():
    dic_dims = [1, 5, 10, 25, 50]

    # Train Koopman model with different number of dictionaries
    dl_solvers = []

    for d in dic_dims:
        solver = FlowDMD(dim=2, num_layers=3, hidden_dim=64,
                         reg=1e-4, num_dictionaries=d)
        dl_solvers.append(solver)
        print("Model has %d trainable parameters." % d)

    n_test = 50
    dl_errors = []
    traj_len = 50  # length of test trajectories

    for i in range(len(dl_solvers)):
        solver = dl_solvers[i]
        dl_error = 0
        for j in range(n_test):
            if j % 10 == 0:
                print("Model no.: ", i, " Test no.: ", j)
            # Generate testing data
            duffing_test = DuffingOscillator(n_init=1, traj_len=traj_len)
            data_x_test = duffing_test.generate_init_data()
            data_x_test = data_x_test.float()

            # Exact trajectory
            x_traj = data_x_test.numpy()

            # Estimated trajectory from DL
            x0_test = data_x_test[0]
            x0_test = x0_test.reshape(-1, x0_test.shape[-1])
            with torch.no_grad():
                x_est_traj_DL = solver.predict(x0_test, traj_len).numpy()

            # 修改误差计算方式，避免数值溢出
            diff = x_est_traj_DL - x_traj
            # 使用更稳定的方式计算误差
            error = np.sqrt(np.mean(np.abs(np.square(x_est_traj_DL-x_traj))))

            if not np.isnan(error):
                dl_error += error
        dl_errors.append(dl_error/n_test)
    rbf_dic_dims = [1, 5, 10, 25, 50]

    rbf_solvers = []
    for d in rbf_dic_dims:
        rbf_basis_func = DicRBF(rbf_number=d, regularizer=1e-4)
        rbf_basis_func.build(data_x_train)
        rbf_solver = KoopmanGeneralSolver(dic=rbf_basis_func,
                                          target_dim=2,
                                          reg=0.0)
        rbf_solver.build(data_train)
        rbf_solvers.append(rbf_solver)

    n_test = 50
    rbf_errors = []
    traj_len = 50  # length of test trajectories

    for i in range(len(rbf_solvers)):
        solver = rbf_solvers[i]
        rbf_error = 0
        for j in range(n_test):
            if j % 10 == 0:
                print("Model no.: ", i, " Test no.: ", j)
            # Generate testing data
            duffing_test = DuffingOscillator(n_init=1, traj_len=traj_len)
            data_x_test = duffing_test.generate_init_data()

            # Exact trajectory
            x_traj = data_x_test.numpy()

            # Estimated trajectory from rbf
            x0_test = data_x_test[0]
            x0_test = x0_test.reshape(-1, x0_test.shape[-1])
            x_est_traj_rbf = solver.predict(x0_test, traj_len).numpy()

            # 修改误差计算方式，避免数值溢出
            diff = x_est_traj_rbf - x_traj
            # 使用更稳定的方式计算误差
            error = np.sqrt(np.mean(np.sum(diff * diff, axis=1)))

            if not np.isnan(error):
                rbf_error += error
        rbf_errors.append(rbf_error/n_test)
    rbf_errors = rbf_errors

    ddims = [1, 5, 10, 25, 50]
    matplotlib.rcParams.update({'font.size': 18})

    plt.plot(ddims, dl_errors, "-o", label="EDMD-Flow")
    plt.plot(ddims, rbf_errors, "-s", label="EDMD (RBF)")
    plt.legend(loc="best")
    plt.ylabel("Reconstruction Error")
    plt.xlabel("Number of dictionary elements")
    plt.savefig('test2.png')
    plt.show()
    plt.close()


if __name__ == "__main__":
    test()
